# state_graph_state_demo.py
# =====================================
# 这个示例演示 LangGraph 中“状态图（StateGraph）”的基本原理，
# 并重点说明“STATE（状态）”在其中扮演的角色。
# =====================================

# LangGraph 提供了一种声明式的“状态流转”机制。
# 你可以把整个流程想象成“节点 + 状态”的图结构：
#   - 每个节点（node）是一个独立的逻辑处理单元。
#   - 状态（state）是一份共享的数据，会在节点之间传递和更新。
#   - StateGraph 则是组织这些节点和状态流转的控制器。

# 导入 START 表示图的起点，StateGraph 是构建器
from langgraph.graph import START, StateGraph
from typing_extensions import TypedDict


# =====================================
# 一、定义状态结构（STATE）
# =====================================
# STATE 是整个状态图中流动的数据结构。
# 每个节点执行时，会接收到当前 state，
# 并返回一个（部分或全部）更新后的新 state。
#
# 我们可以用 Python 的 TypedDict 来定义 state 的字段，
# 这样 IDE 能帮我们提示和检查。
class MyState(TypedDict):
    x: int
    y: int
    sum: int  # 额外字段，用于保存计算结果


# =====================================
# 二、定义节点函数（NODE）
# =====================================
# 每个节点都是一个普通的 Python 函数。
# 它必须接收 state（当前状态）和可选的 config，
# 并返回一个新的字典，用于更新全局 state。

def add_node(state: MyState, config=None) -> MyState:
    """
    节点1：add_node
    功能：对 x 和 y 做简单加法操作。
    """
    new_x = state["x"] + 1
    new_y = state["y"] + 2
    print(f"[add_node] 输入状态: {state} → 输出状态: x={new_x}, y={new_y}")
    return {"x": new_x, "y": new_y, "sum": state["sum"]}  # 保留 sum


def sum_node(state: MyState, config=None) -> MyState:
    """
    节点2：sum_node
    功能：计算 x + y，并写入到 sum 字段中。
    """
    total = state["x"] + state["y"]
    print(f"[sum_node] 接收状态: {state} → 计算结果 sum={total}")
    return {"x": state["x"], "y": state["y"], "sum": total}


# =====================================
# 三、构建状态图（StateGraph）
# =====================================
# 通过 builder 定义节点之间的连接关系。

builder = StateGraph(MyState)

# 添加两个节点
builder.add_node(add_node)
builder.add_node(sum_node)

# 添加状态流转关系（edges）
# START → add_node → sum_node
builder.add_edge(START, "add_node")
builder.add_edge("add_node", "sum_node")

# 编译为可执行的状态图对象
graph = builder.compile()

# =====================================
# 四、执行状态图
# =====================================
# 初始状态（STATE）包含所有字段的初始值。
# 这相当于整个流程的“输入上下文”。
initial_state = {"x": 1, "y": 2, "sum": 0}

print("\n=== 状态图执行开始 ===")
# invoke() 会自动根据定义的流转顺序执行节点，
# 每个节点执行后，都会返回更新的状态，
# 并把新的状态传递给下一个节点。
final_state = graph.invoke(initial_state)
print("=== 状态图执行结束 ===\n")

print("最终输出状态：", final_state)
